import json
import pandas as pd
import mplfinance as mpf
import matplotlib.pyplot as plt
import chardet
import numpy as np
from matplotlib import font_manager as fm
from matplotlib.dates import DateFormatter, DayLocator
import matplotlib as mpl
import matplotlib.ticker as mticker
import matplotlib.dates as mdates

# 设置全局样式 - 匹配图片
plt.rcParams['font.sans-serif'] = ['SimHei']  # 中文字体
plt.rcParams['axes.unicode_minus'] = False   # 解决负号显示问题
plt.rcParams['axes.grid'] = True             # 显示网格
plt.rcParams['grid.linestyle'] = ':'         # 网格线样式
plt.rcParams['grid.alpha'] = 0.3             # 网格线透明度
mpl.rcParams['figure.dpi'] = 200             # 调整分辨率

# 修正刻度问题 - 关键修复
mpl.rcParams['axes.autolimit_mode'] = 'round_numbers'
mpl.rcParams['axes.xmargin'] = 0.01
mpl.rcParams['axes.ymargin'] = 0.05

# 定义颜色方案 - 匹配图片
BG_COLOR = '#FFFFFF'    # 白色背景
GRID_COLOR = '#C0C0C0'  # 浅灰色网格
BUY_COLOR = '#008000'   # 绿色买入标记
SELL_COLOR = '#FF0000'  # 红色卖出标记
BUY_LABEL_COLOR = '#90EE90'  # 买入标签背景色
SELL_LABEL_COLOR = '#FFC0CB'  # 卖出标签背景色

start_date = '20250529'
end_date = '20250714'
kline_trade_file = 'D:\\work\\code\\clark\\gitee\\big_a\\datas\\backtest\\sz002468\\20250529-20250714_chart.txt'

def detect_encoding(file_path):
    with open(file_path, 'rb') as f:
        return chardet.detect(f.read(10000))['encoding']

def load_kline_data(file_path):
    kline_data = []
    encoding = detect_encoding(file_path)
    with open(file_path, 'r', encoding=encoding, errors='replace') as file:
        for line in file:
            if line.strip().startswith('{"add":'):
                try:
                    kline_data.append(json.loads(line))
                except json.JSONDecodeError:
                    continue
    return kline_data

def load_trade_data(file_path):
    trade_data = []
    encoding = detect_encoding(file_path)
    with open(file_path, 'r', encoding=encoding, errors='replace') as file:
        for line in file:
            if line.strip().startswith('{"cash":'):
                try:
                    trade_data.append(json.loads(line))
                except json.JSONDecodeError:
                    continue
    return trade_data

def main():
    print("开始处理数据...")
    try:
        # 加载数据
        kline_data = load_kline_data(kline_trade_file)
        if not kline_data:
            raise ValueError("未加载任何K线数据")
        print(f"加载了 {len(kline_data)} 条K线数据")

        trade_data = load_trade_data(kline_trade_file)
        if not trade_data:
            print("警告: 未加载任何交易数据")
        else:
            print(f"加载了 {len(trade_data)} 条交易数据")

        # 处理K线数据
        df = pd.DataFrame(kline_data)
        if 'date' not in df.columns:
            raise ValueError("K线数据缺少'date'列")

        df['date'] = pd.to_datetime(df['date'])
        df.set_index('date', inplace=True)

        # 创建OHLC数据 - 匹配列名
        ohlc = df.rename(columns={
            'open': 'Open',
            'high': 'High',
            'low': 'Low',
            'close': 'Close',
            'volume': 'Volume'
        })[['Open', 'High', 'Low', 'Close', 'Volume']]

        # 确保数据类型正确
        ohlc = ohlc.apply(pd.to_numeric, errors='coerce')

        # 处理交易数据
        trades = pd.DataFrame(trade_data)
        if trades.empty:
            print("警告: 未找到买卖点数据")
            buy_points = pd.DataFrame(columns=['date', 'price', 'volume'])
            sell_points = pd.DataFrame(columns=['date', 'price', 'volume'])
        else:
            # 确保有必要的列
            if 'date' not in trades.columns:
                trades['date'] = pd.to_datetime('now')

            trades['date'] = pd.to_datetime(trades['date'])
            trades = trades[trades['operation'].isin(['买入', '卖出'])]

            # 准备买卖点数据
            buy_points = trades[trades['operation'] == '买入'][['date', 'price', 'volume']].copy()
            sell_points = trades[trades['operation'] == '卖出'][['date', 'price', 'volume']].copy()

            # 确保数值类型
            buy_points['price'] = buy_points['price'].astype(float)
            buy_points['volume'] = buy_points['volume'].astype(int)
            sell_points['price'] = sell_points['price'].astype(float)
            sell_points['volume'] = sell_points['volume'].astype(int)

        print(f"K线数据时间范围: {ohlc.index.min()} 至 {ohlc.index.max()}")
        print(f"共找到 {len(buy_points) + len(sell_points)} 个交易点")

        # 1. 创建图表框架 - 匹配图片布局
        fig, axes = plt.subplots(2, 1, figsize=(15, 10), gridspec_kw={'height_ratios': [3, 1]}, sharex=True)
        ax1 = axes[0]  # 价格图
        ax2 = axes[1]  # 成交量图

        # 2. 绘制K线图 - 解决刻度过多问题
        # 关键修复1: 设置X轴限制防止过多刻度
        date_range = ohlc.index.max() - ohlc.index.min()
        if date_range.days > 30:
            # 对长时间数据使用样本
            ohlc_sampled = ohlc.resample('30T').agg({
                'Open': 'first',
                'High': 'max',
                'Low': 'min',
                'Close': 'last',
                'Volume': 'sum'
            })
            print(f"数据量过大({len(ohlc)}条)，采样至{len(ohlc_sampled)}条")
        else:
            ohlc_sampled = ohlc

        # 关键修复2: 使用mpf.plot但避免刻度问题
        mpf.plot(ohlc_sampled,
                 type='candle',
                 ax=ax1,
                 volume=ax2,
                 style='yahoo',
                 show_nontrading=False,
                 update_width_config=dict(candle_width=0.6),
                 warn_too_much_data=len(ohlc_sampled)+1)  # 避免过多数据警告

        # 3. 添加买卖点标记 - 匹配图片
        if not buy_points.empty:
            # 匹配图片中的绿色上三角形
            ax1.plot(buy_points['date'], buy_points['price'], 'g^',
                     markersize=12,  # 匹配图片尺寸
                     markerfacecolor=BUY_COLOR,
                     markeredgecolor='black',
                     linewidth=0.5,
                     label='买入',
                     zorder=10)

        if not sell_points.empty:
            # 匹配图片中的红色下三角形
            ax1.plot(sell_points['date'], sell_points['price'], 'rv',
                     markersize=12,  # 匹配图片尺寸
                     markerfacecolor=SELL_COLOR,
                     markeredgecolor='black',
                     linewidth=0.5,
                     label='卖出',
                     zorder=10)

        # 4. 添加交易标注文本 - 完全匹配图片中的格式
        # 示例图片中的文本: "价格:11.41\n数量:8900"
        def format_trade_info(price, volume):
            """匹配图片中的文本格式"""
            return f"价格:{price:.2f}\n数量:{volume}"

        # 买入点标注 - 匹配图片中上方位置
        for _, row in buy_points.iterrows():
            ax1.annotate(
                format_trade_info(row['price'], row['volume']),
                xy=(row['date'], row['price']),
                xytext=(0, 25),  # 匹配图片中的上方偏移
                textcoords='offset points',
                ha='center',
                va='bottom',
                fontsize=10,
                fontweight='bold',
                bbox=dict(
                    boxstyle='round,pad=0.5',
                    facecolor=BUY_LABEL_COLOR,
                    alpha=0.95,
                    edgecolor=BUY_COLOR,
                    linewidth=1
                ),
                color='black',
                zorder=15
            )

        # 卖出点标注 - 匹配图片中下方位置
        for _, row in sell_points.iterrows():
            ax1.annotate(
                format_trade_info(row['price'], row['volume']),
                xy=(row['date'], row['price']),
                xytext=(0, -30),  # 匹配图片中的下方偏移
                textcoords='offset points',
                ha='center',
                va='top',
                fontsize=10,
                fontweight='bold',
                bbox=dict(
                    boxstyle='round,pad=0.5',
                    facecolor=SELL_LABEL_COLOR,
                    alpha=0.95,
                    edgecolor=SELL_COLOR,
                    linewidth=1
                ),
                color='black',
                zorder=15
            )

        # 5. 设置标题和标签 - 匹配图片
        title_text = f'资产价格走势与交易点 ({start_date} 至 {end_date})'
        ax1.set_title(title_text, fontsize=14, fontweight='bold', pad=15)
        ax1.set_ylabel('价格', fontsize=12, labelpad=10)
        ax2.set_ylabel('成交量', fontsize=10, labelpad=10)

        # 6. 设置X轴格式 - 关键修复解决刻度过多问题
        # 关键修复3: 使用DayLocator解决刻度过多问题
        ax1.xaxis.set_major_locator(DayLocator(interval=7))  # 每7天一个主刻度
        ax1.xaxis.set_major_formatter(DateFormatter('%m-%d'))  # 匹配图片中的日期格式

        # 关键修复4: 手动设置X轴范围
        date_min = ohlc_sampled.index.min() - pd.Timedelta(hours=6)
        date_max = ohlc_sampled.index.max() + pd.Timedelta(hours=6)
        ax1.set_xlim(date_min, date_max)

        # 7. 设置Y轴格式
        ax1.yaxis.set_major_formatter(mticker.FormatStrFormatter('%.2f'))

        # 8. 添加图例 - 匹配图片中的位置
        ax1.legend(loc='upper right', fontsize=10, framealpha=0.8)

        # 9. 调整布局
        plt.subplots_adjust(top=0.94, bottom=0.08, left=0.06, right=0.98, hspace=0.1)

        # 保存结果
        output_file = f'资产价格走势图_{start_date}_{end_date}.png'
        plt.savefig(output_file, dpi=200, bbox_inches='tight')
        print(f"图表已保存为: {output_file}")

        plt.show()

    except Exception as e:
        print(f"处理过程中发生错误: {str(e)}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    main()